/*
 *  Copyright 2013 National Institute of Advanced Industrial Science and Technology
 *  
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *  
 *      http://www.apache.org/licenses/LICENSE-2.0
 *  
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

package org.sss.dummykvs;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.SortedMap;
import java.util.TreeMap;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.sss.mapreduce.DB;
import org.sss.util.BytesComparator;

public class TreeMapDBBackend implements DB {
  private static Logger logger = LoggerFactory.getLogger(TreeMapDBBackend.class);

  private BytesComparator cmp = new BytesComparator();
  private TreeMap<byte[], byte[]> map = new TreeMap<byte[], byte[]>(cmp);

  public synchronized int size() {
    return map.size();
  }

  @Override
  public synchronized void putlist(byte[][] kvs) throws IOException {
    for (int i = 0; i < kvs.length; i+=2) {
      map.put(kvs[i+0], kvs[i+1]);
    }
  }

  @Override
  public synchronized List<byte[]> range(byte[] start, byte[] excludingEnd, int limit,
                                         boolean includeStart) throws IOException {
    List<byte[]> tmp  = new ArrayList<byte[]>();

    int size = 0;
    if (limit <= 0) {
      return tmp;
    }
    if (includeStart) {
      byte [] exact = map.get(start);
      if (exact != null) {
        tmp.add(start);
        tmp.add(exact);
        size += start.length + exact.length;
      }
    }
    byte [] last = start;
    while (size < limit) {
      Map.Entry<byte[], byte[]> entry = map.higherEntry(last);
      if (entry == null) break;
      if (cmp.compare(entry.getKey(), excludingEnd) >= 0)
        break;
      tmp.add(entry.getKey());
      tmp.add(entry.getValue());
      size += entry.getKey().length + entry.getValue().length;
      last = entry.getKey();
    }
    return tmp;
  }

  @Override
  public synchronized List<byte[]> ranged(byte[] start, byte[] excludingEnd, int limit,
                                          boolean includeStart) throws IOException {
    List<byte[]> tmp  = new ArrayList<byte[]>();

    int size = 0;
    if (limit <= 0) {
      return tmp;
    }
    if (includeStart) {
      byte [] exact = map.get(start);
      if (exact != null) {
        tmp.add(start);
        tmp.add(exact);
        size += start.length + exact.length;
        map.remove(start);
      }
    }
    byte [] last = start;
    while (size < limit) {
      Map.Entry<byte[], byte[]> entry = map.higherEntry(last);
      if (entry == null) break;
      if (cmp.compare(entry.getKey(), excludingEnd) >= 0)
        break;
      tmp.add(entry.getKey());
      tmp.add(entry.getValue());
      size += entry.getKey().length + entry.getValue().length;
      last = entry.getKey();
      map.remove(entry.getKey());
    }
    return tmp;
  }

  @Override
  public synchronized int rangeout(byte[] start, byte[] excludingEnd, int maxCount,
                                   boolean includeStart) throws IOException {
    if (maxCount < 0) {
      maxCount = map.size();
    }
    int count = 0;
    for (; count < maxCount;++count) {
      SortedMap<byte[], byte[]> s = map.subMap(start, includeStart, excludingEnd, false);

      if (s.isEmpty()) {
        break;
      }
      map.remove(s.firstKey());
    }
    return count;
  }

  @Override
  public synchronized List<byte[]> fwmkeys(byte[] prefix, int max) throws IOException {
    byte[] tail = new byte[prefix.length];
    int i = tail.length;
    int carry = 1;
    while (i > 0) {
      i--;
      tail[i] = (byte)(prefix[i] + carry);
      carry = (tail[i] == 0)?carry:0;
    }

    SortedMap<byte[], byte[]> m = null;
    if (carry == 0) {
      logger.debug("submap from {} to {}", Arrays.toString(prefix), Arrays.toString(tail));
      m = map.subMap(prefix, tail);
    } else {
      logger.debug("submap from {}", Arrays.toString(prefix));
      m = map.tailMap(prefix);
    }

    List<byte[]> tmp  = new ArrayList<byte[]>();

    if (max < 0) {
      max = m.size();
    }
    int n = 0;
    for (byte[] b: m.keySet()) {
      if (n++ >= max) {
        break;
      }

      tmp.add(b);
    }

    return tmp;
  }

  @Override
  public synchronized byte[] get(byte[] key) throws IOException {
    return map.get(key);
  }

  @Override
  public void close() throws IOException {
    // Do nothing
  }

  @Override
  public synchronized List<byte[]> getlist(byte[][] keys) throws IOException {
    List<byte[]> tmp = new ArrayList<byte[]>();
    for (byte[] key: keys) {
      byte[] val = map.get(key);
      if (val != null) {
        tmp.add(key);
        tmp.add(val);
      }
    }
    return tmp;
  }

  @Override
  public void outlist(byte[][] keys) throws IOException {
    for (byte[] key: keys) {
      map.remove(key);
    }
  }

  // For test
  public void put(String k, String v) throws IOException {
    put(k.getBytes("UTF-8"), v.getBytes("UTF-8"));
  }

  public void put(byte[] k, byte[] v) throws IOException {
    byte[][] bytes =  new byte[2][];
    bytes[0] = k;
    bytes[1] = v;
    putlist(bytes);
  }

  public static void main(String [] str) throws IOException {
    TreeMapDBBackend db = new TreeMapDBBackend();
    db.put("aa0a", "a");
    db.put("aaaa", "a");
    db.put("aaab", "b");
    db.put("aaac", "c");
    db.put("aaacd", "c");
    db.put("aaaed", "c");
    db.put("aabed", "c");
    for (byte[] bytes: db.range("aaab".getBytes(), "bbbb".getBytes(), 10, true))
      System.out.println(new String(bytes));
  }
}
